import pinocchio as pin
from pinocchio.utils import rand, zero
import meshcat

# 引入src中所有的函数
import sys
from os.path import dirname, join

sys.path.append(join(dirname(dirname(__file__)), "src"))
sys.path.append(dirname(dirname(__file__))) # 把文件根路径添加到sys.path中，这样在根路径下的src文件夹才能被import
from src import *
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 获取URDF模型路径
URDF_FILE_PATH = filePath(URDF_NAME = 'ur5_gripper', ROBOT = True)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 构建robot对象，并得到viz
robot, viz = robot_viz(URDF_FILE_PATH)
q0 = pin.neutral(robot.model)
pin.framesForwardKinematics(robot.model, robot.data, q0)
pin.updateFramePlacements(robot.model, robot.data) 
robot.initViewer(loadModel=True)
viz.display(q0)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


model = robot.model
data = robot.data

#####################################################
#####################################################
#####################################################
#####################################################
#####################################################

# end effector 是第六个关节
JOINT_ID = 6

# 期望的姿态如下（SE3）
oMdes = pin.SE3(np.eye(3), np.array([0.5, 0.0, 0.0]))

# 设置初始参数
q = pin.neutral(model)

# 设定部分计算相关常量
eps = 1e-4
IT_MAX = 1000
DT = 1e-2
damp = 1e-12

i = 0
while True:
    pin.forwardKinematics(model, data, q)
    # dMi corresponds to the transformation between the desired pose and the current one
    iMd = data.oMi[JOINT_ID].actInv(oMdes)

    # 通过对数映射将李群SE(3)空间中的4x4变换矩阵映射为一个六维的se(3)李代数向量，再求这个转动向量的范数
    err = pin.log(iMd).vector # in joint frame
    if norm(err) < eps:
        success = True
        break
    if i >= IT_MAX:
        success = False
        break
    print(f"iteration: {i}, error: {err}")
    # 计算基坐标到JOINT_ID的Jacobian
    J = pin.computeJointJacobian(model, data, q, JOINT_ID) # in joint frame
    # 得到IK需要的Jacobian，即表示error的Jacobian，https://scaron.info/robotics/jacobian-of-a-kinematic-task-and-derivatives-on-manifolds.html
    J = -np.dot(pin.Jlog6(iMd.inverse()), J)
    
    # 为了避免出现奇异，这里先计算error Jacobian的damped pseudo-inverse
    v = -J.T.dot(solve(J.dot(J.T) + damp * np.eye(6), err))
    # 然后基于error Jacobian的伪逆（v），去迭代q
    q = pin.integrate(model, q, v * DT)
    viz.display(q)


if success:
    print("Convergence achieved!")
else:
    print(
        "\n"
        "Warning: the iterative algorithm has not reached convergence "
        "to the desired precision"
    )

print(f"\nresult: {q.flatten().tolist()}")
print(f"\nfinal error: {err.T}")